import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

test_data = torchvision.datasets.CIFAR10(
    # 数据集下载路径
    "./dataset",
    download=True,
    # 测试数据集还是训练数据集
    train=False,
    # 图片是PIL格式，需要转换为tensor
    transform=torchvision.transforms.ToTensor()
)

test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True,num_workers=0,drop_last=False)

# 第一张图片及target
img, target = test_data[0]
print(img.shape, target)

writer = SummaryWriter("../logs")
step = 0
for data in test_loader:
    imgs, targets = data
    # print(imgs.shape, targets)
    writer.add_images("test_data", imgs, step)
    step += 1

writer.close()